#!/usr/bin/env python3
import z3
from z3 import *
import unittest


class TestSimpleEquation(unittest.TestCase):

    def test_solve_equation(self):
        """
        x + 2y = 7 where x > 2 and y < 10
        """
        x, y = z3.Ints('x y')

        solver = z3.Solver()

        solver.add(x + 2 * y == 7)
        solver.add(x > 2)
        solver.add(y == 9)

        if solver.check() == z3.sat:
            print(solver.model())
        else:
            print("no result")

    def test_solve_equation2(self):
        # 上面的例子也可以这样写
        x = Int('x')
        y = Int('y')
        solve(x + 2 * y == 7, x > 2, y == 0)

    def test_simplify_equation(self):
        x = Int('x')
        y = Int('y')
        # html_mode会让结果以html标签格式显示
        # set_option(html_mode=True)
        print(simplify(x + y + 2 * x + 3))
        print(simplify(x < y + x + 2))
        print(simplify(And(x + 1 >= 3, x ** 2 + x ** 2 + y ** 2 + 2 >= 5)))

    def test_solve_real(self):
        x = Real('x')
        y = Real('y')
        solve(x ** 2 + y ** 2 > 3)

    def test_solve_real_2(self):
        a, b, c, d = Reals('a b c d')
        solve(
            a + b == -1,
            c + d == 2,
            3 * a + 7 * b == 2,
            3 * c + 7 * d == 1
        )
